package contest.oop.sse.person;

import contest.oop.sse.key.KeyLoader;
import com.sse.iitp.message.*;
import contest.oop.sse.message.*;
import contest.oop.sse.secalg.CaesarAlgorithm;
import contest.oop.sse.secalg.SymmetricAlgorithm;
import contest.oop.sse.secalg.XorAlgorithm;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.util.ArrayList;
import java.util.Base64;
import java.util.Random;

/**
 * 共产党地下工作者类
 *
 * @author jwmao
 */
public class CpcAgent extends Agent {
    private PrivateKey privateKey;
    private PublicKey publicKey;
    private int[] randomNumbers;
    private SymmetricAlgorithm algorithm;
    private String algorithmName;
    private static final String OK = "OK";

    public CpcAgent(String name) {
        this.name = name;
        randomNumbers = new int[]{-1, -1, -1};
        this.connectors = new ArrayList<>();
    }

    public void getPrivateKey(KeyLoader loader, String keyId) {
        this.privateKey = loader.obtainPrivateKey(keyId);
    }

    public void getPublicKey(KeyLoader loader, String keyId) {
        this.publicKey = loader.obtainPublicKey(keyId);
    }

    public void setSymmetricAlgorithm(SymmetricAlgorithm algorithm) {
        this.algorithm = algorithm;
    }

    @Override
    public void receive(Message message) {
        this.message = message;
        System.out.println("地下工作者" + name + " 收到原始消息：" + message.getBody());
        processMessage();
    }

    private void processMessage() {
        if (this.message instanceof ClientHelloMessage) {
            this.randomNumbers[0] = ((ClientHelloMessage) this.message).getRandom1();
            String[] availableAlgorithms = ((ClientHelloMessage) this.message).getAvailableAlgorithmList().split(",");
            double rand = new Random().nextDouble();
            if (rand > 0.5) {
                algorithmName = availableAlgorithms[1];
            } else {
                algorithmName = availableAlgorithms[0];
            }
        } else if (this.message instanceof ServerHelloMessage) {
            this.randomNumbers[1] = ((ServerHelloMessage) this.message).getRandom2();
            this.algorithmName = ((ServerHelloMessage) this.message).getChosenAlgorithm();
            this.publicKey = ((ServerHelloMessage) this.message).getPublicKey();
            if ("Caesar".equals(this.algorithmName)) {
                this.algorithm = new CaesarAlgorithm();
            } else if ("Xor".equals(this.algorithmName)) {
                this.algorithm = new XorAlgorithm();
            }
        } else if (this.message instanceof HandshakeMessage) {
            if (this.randomNumbers[2] == -1) {
                this.randomNumbers[2] = Integer.parseInt(rsaDecrypt(this.message.getBody(), privateKey));
                int key = this.randomNumbers[0] + this.randomNumbers[1] + this.randomNumbers[2];
                if (this.algorithm instanceof CaesarAlgorithm) {
                    this.algorithm.setKey(Integer.toString(key % 26));
                } else if (this.algorithm instanceof XorAlgorithm) {
                    this.algorithm.setKey(Integer.toString(key & 0xFF));
                }
            } else if (OK.equals(this.algorithm.decrypt(Base64.getDecoder().decode(this.message.getBody())))) {
                System.out.println("地下工作者" + name + "，已建立联系。");
            } else {
                System.out.println("消息发送错误，退出。");
                System.exit(-1);
            }
        } else {
            if (this.algorithm != null) {
                String messageBody = this.algorithm.decrypt(Base64.getDecoder().decode(this.message.getBody()));
                System.out.println("地下工作者" + name + " 收到：" + messageBody);
            }
            else {
                System.out.println("错误！对称加密算法未设置！");
                System.exit(-1);
            }
        }
    }

    private String rsaDecrypt(String body, PrivateKey privateKey) {
        byte[] inputByte = Base64.getDecoder().decode(body.getBytes(StandardCharsets.UTF_8));
        try {
            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
            return new String(cipher.doFinal(inputByte));
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) {
            e.printStackTrace();
        }
        return "";
    }

    private String rsaEncrypt(String body, PublicKey publicKey) {
        Cipher cipher;
        try {
            cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
            return Base64.getEncoder().encodeToString(cipher.doFinal(body.getBytes(StandardCharsets.UTF_8)));
        } catch (NoSuchPaddingException | IllegalBlockSizeException | NoSuchAlgorithmException | BadPaddingException | InvalidKeyException e) {
            e.printStackTrace();
        }
        return "";
    }

    @Override
    public void createMessage(MessageType messageType) {
        switch (messageType) {
            case CLIENT_HELLO:
                this.randomNumbers[0] = new Random().nextInt(64);
                String availableAlgorithms = "Caesar,Xor";
                this.message = new ClientHelloMessage(this.randomNumbers[0], availableAlgorithms);
                break;
            case SERVER_HELLO:
                this.randomNumbers[1] = new Random().nextInt(64);
                if ("Caesar".equals(this.algorithmName)) {
                    this.algorithm = new CaesarAlgorithm();
                } else if ("Xor".equals(this.algorithmName)) {
                    this.algorithm = new XorAlgorithm();
                }
                this.message = new ServerHelloMessage(this.randomNumbers[1], this.algorithmName, this.publicKey);
                break;
            case HANDSHAKE:
                if (this.randomNumbers[2] == -1) {
                    this.randomNumbers[2] = new Random().nextInt(64);
                    int key = this.randomNumbers[0] + this.randomNumbers[1] + this.randomNumbers[2];
                    if (this.algorithm instanceof CaesarAlgorithm) {
                        this.algorithm.setKey(Integer.toString(key % 26));
                    } else if (this.algorithm instanceof XorAlgorithm) {
                        this.algorithm.setKey(Integer.toString(key & 0xFF));
                    }
                    String body = rsaEncrypt(Integer.toString(this.randomNumbers[2]), publicKey);
                    createMessage(messageType, body);
                }
                else {
                    String body = Base64.getEncoder().encodeToString(this.algorithm.encrypt("OK"));
                    createMessage(messageType, body);
                }
                break;
            default:
                break;
        }
    }

    @Override
    public void createMessage(MessageType messageType, String msgBody) {
        switch (messageType) {
            case CLIENT_HELLO:
                this.randomNumbers[0] = new Random().nextInt(64);
                String availableAlgorithms = "Caesar,Xor";
                this.message = new ClientHelloMessage(this.randomNumbers[0], availableAlgorithms);
                this.message.setBody(msgBody);
                break;
            case SERVER_HELLO:
                this.randomNumbers[1] = new Random().nextInt(64);
                if ("Caesar".equals(this.algorithmName)) {
                    this.algorithm = new CaesarAlgorithm();
                } else if ("Xor".equals(this.algorithmName)) {
                    this.algorithm = new XorAlgorithm();
                }
                this.message = new ServerHelloMessage(this.randomNumbers[1], this.algorithmName, this.publicKey);
                this.message.setBody(msgBody);
                break;
            case HANDSHAKE:
                this.message = new HandshakeMessage(msgBody);
                break;
            case STANDARD:
                String body = Base64.getEncoder().encodeToString(this.algorithm.encrypt(msgBody));
                this.message = new StandardMessage(body);
                break;
            default:
                break;
        }
    }
}
